import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import sklearn
import sys
import os
import pandas as pd
import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from PIL import Image

#
# train_covid_image_file = '../dataset/Data-split/COVID/trainCT_COVID.txt'
# test_covid_image_file='../dataset/Data-split/COVID/testCT_COVID.txt'
# val_covid_image_file='../dataset/Data-split/COVID/valCT_COVID.txt'

# train_Nocovid_image_file = '../dataset/Data-split/NonCOVID/trainCT_NonCOVID.txt'
# test_Nocovid_image_file='../dataset/Data-split/NonCOVID/testCT_NonCOVID.txt'
# val_Nocovid_image_file='../dataset/Data-split/NonCOVID/valCT_NonCOVID.txt'

root_directory_COVID = "../dataset/data4/train/COVID/"
root_directory_NoCOVID = "../dataset/data4/train/NOCOVID/"


def parse_directory_file(directory_Path_Name):
    results = []
    for filename in os.listdir(r"./" + directory_Path_Name):
        results.append(directory_Path_Name + filename)
    return results


covid_image_list_COVID = parse_directory_file(root_directory_COVID)
Nocovid_image_list_COVID = parse_directory_file(root_directory_NoCOVID)

print("covid_image_list_COVID:", len(covid_image_list_COVID))
print("Nocovid_image_list_COVID", len(Nocovid_image_list_COVID))
rar = 0.0
print("--------------生成COVID------------------")
for image in covid_image_list_COVID:
    datagen = ImageDataGenerator(
        rotation_range=15 + rar,
        width_shift_range=0.08,
        height_shift_range=0.08,
        shear_range=0.1,
        zoom_range=0.1,
        fill_mode='nearest')
    rar = rar + 0.5
    img = load_img(image)  # 这是一个PIL图像
    x = img_to_array(img)
    x = x.reshape((1,) + x.shape)
    i = 0
    for batch in datagen.flow(x, batch_size=1,
                              save_to_dir=root_directory_COVID,
                              save_prefix='plate',
                              save_format='png'):
        i += 1
        print("<" + image + "图像> 生成第" + str(i) + " 张图像")
        if i > 3:
            break
    print("----------------------")

rar = 0.0
print("--------------生成NOCOVID------------------")
for image in Nocovid_image_list_COVID:
    datagen = ImageDataGenerator(
        rotation_range=15 + rar,
        width_shift_range=0.08,
        height_shift_range=0.08,
        shear_range=0.1,
        zoom_range=0.1,
        fill_mode='nearest')
    rar = rar + 0.5
    img = load_img(image)  # 这是一个PIL图像
    x = img_to_array(img)
    x = x.reshape((1,) + x.shape)
    i = 0
    for batch in datagen.flow(x, batch_size=1,
                              save_to_dir=root_directory_NoCOVID, save_prefix='plate', save_format='png'):
        i += 1
        print("<" + image + "图像> 生成第" + str(i) + " 张图像")
        if i > 3:
            break
    print("----------------------")
print("=================增强完毕=====================")

covid_image_list_COVID = parse_directory_file(root_directory_COVID)
Nocovid_image_list_COVID = parse_directory_file(root_directory_NoCOVID)

print("covid_image_list_COVID:", len(covid_image_list_COVID))
print("Nocovid_image_list_COVID", len(Nocovid_image_list_COVID))